#include "rknn_api.h"
#include <cstring>
#include <iomanip>
#include <iostream>
#include <vector>

static unsigned char *load_model(const char *filename, int *model_size)
{
    FILE *fp;
    unsigned char *data;

    fp = fopen(filename, "rb");
    if (NULL == fp)
    {
        printf("Open file %s failed.\n", filename);
        return NULL;
    }

    fseek(fp, 0, SEEK_END);
    int size = ftell(fp);

    data = (unsigned char *)malloc(size);
    if (data == NULL)
    {
        printf("buffer malloc failure.\n");
        fclose(fp);
        return NULL;
    }

    fseek(fp, 0, SEEK_SET);
    fread(data, 1, size, fp);
    fclose(fp);

    *model_size = size;
    return data;
}

int main()
{
    const char *model_path = "model/wuzi_640x640.rknn";
    int model_data_size = 0;
    unsigned char *model_data = load_model(model_path, &model_data_size);

    if (model_data == NULL)
    {
        printf("Failed to load model file: %s\n", model_path);
        return -1;
    }

    printf("Model loaded successfully, size: %d bytes\n", model_data_size);

    rknn_context ctx;
    int ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
    if (ret < 0)
    {
        printf("rknn_init error ret=%d\n", ret);
        free(model_data);
        return -1;
    }

    // 查询输入输出数量
    rknn_input_output_num io_num;
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret < 0)
    {
        printf("rknn_query IN_OUT_NUM error ret=%d\n", ret);
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }
    printf("Model has %d inputs, %d outputs\n", io_num.n_input, io_num.n_output);

    // 查询输入属性
    rknn_tensor_attr input_attr;
    input_attr.index = 0;
    ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &input_attr, sizeof(rknn_tensor_attr));
    if (ret < 0)
    {
        printf("rknn_query INPUT_ATTR error ret=%d\n", ret);
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }

    printf("Input dimensions: [");
    for (uint32_t j = 0; j < input_attr.n_dims; j++)
    {
        printf("%d", input_attr.dims[j]);
        if (j < input_attr.n_dims - 1)
            printf(", ");
    }
    printf("]\n");

    // 查询输出属性
    rknn_tensor_attr output_attr;
    output_attr.index = 0;
    ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &output_attr, sizeof(rknn_tensor_attr));
    if (ret < 0)
    {
        printf("rknn_query OUTPUT_ATTR error ret=%d\n", ret);
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }

    printf("Output dimensions: [");
    for (uint32_t j = 0; j < output_attr.n_dims; j++)
    {
        printf("%d", output_attr.dims[j]);
        if (j < output_attr.n_dims - 1)
            printf(", ");
    }
    printf("]\n");
    printf("Output scale: %f, zp: %d\n", output_attr.scale, output_attr.zp);

    // 创建测试输入数据
    int input_size = input_attr.dims[1] * input_attr.dims[2] * input_attr.dims[3];
    unsigned char *input_data = (unsigned char *)malloc(input_size);
    if (input_data == NULL)
    {
        printf("Failed to allocate input data\n");
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }

    // 填充测试数据（全零）
    memset(input_data, 0, input_size);

    // 设置输入
    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = input_size;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = input_data;

    ret = rknn_inputs_set(ctx, io_num.n_input, inputs);
    if (ret < 0)
    {
        printf("rknn_inputs_set error ret=%d\n", ret);
        free(input_data);
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }

    // 运行推理
    ret = rknn_run(ctx, NULL);
    if (ret < 0)
    {
        printf("rknn_run error ret=%d\n", ret);
        free(input_data);
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }

    // 获取输出
    rknn_output outputs[1];
    memset(outputs, 0, sizeof(outputs));
    outputs[0].want_float = 0; // 获取量化数据

    ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
    if (ret < 0)
    {
        printf("rknn_outputs_get error ret=%d\n", ret);
        free(input_data);
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }

    // 分析输出数据
    int8_t *output_data = (int8_t *)outputs[0].buf;
    int output_size = outputs[0].size;

    printf("Output size: %d bytes\n", output_size);
    printf("Expected size: %d bytes\n", output_attr.size);

    // 打印前几个检测框的数据
    int num_boxes = output_attr.dims[2]; // 6300
    int box_size = output_attr.dims[1];  // 19

    printf("Number of boxes: %d\n", num_boxes);
    printf("Box size: %d\n", box_size);

    // 打印前3个检测框的数据
    for (int i = 0; i < 3 && i < num_boxes; i++)
    {
        printf("Box %d: ", i);
        for (int j = 0; j < box_size; j++)
        {
            int idx = i * box_size + j;
            if (idx < output_size)
            {
                printf("%d ", output_data[idx]);
            }
        }
        printf("\n");
    }

    // 释放资源
    rknn_outputs_release(ctx, io_num.n_output, outputs);
    free(input_data);
    rknn_destroy(ctx);
    free(model_data);

    return 0;
}